#!/usr/bin/env python3
"""
Prompt Manager for GSI/EnKF Automated Code Commentary System

This module handles prompt management, template loading, and comment standards.

Author: GSI Development Team
Date: 2025-01-XX
"""

import logging
from pathlib import Path
from typing import Dict, List, Optional, Any

from ..core.models import CommentStandard, PromptConfig

logger = logging.getLogger(__name__)


class PromptManager:
    """Manages prompts and comment standards with configurable templates"""
    
    def __init__(self, config_dir: str):
        self.config_dir = Path(config_dir)
        self.comment_standard = self._load_comment_standard()
        self.classification_templates = self._load_classification_templates()
        self.custom_prompts = self._load_custom_prompts()
        
    def _load_comment_standard(self) -> CommentStandard:
        """Load structured comment standard with improved templates"""
        return CommentStandard(
            subroutine_template="""!> @brief [Brief description of subroutine purpose]
!> @details [Detailed description of algorithm and implementation approach]
!> [Include mathematical formulation if applicable]
!> 
!> @param[in] [param_name] [Description with units and valid ranges]
!> @param[out] [param_name] [Description with units and expected format] 
!> @param[inout] [param_name] [Description with modification details]
!>
!> @usage
!> call subroutine_name(param1, param2, param3)
!> ! [Brief example showing typical usage patterns]
!>
!> @algorithm  
!> 1. [Step 1: Initial setup and validation]
!> 2. [Step 2: Main computation or processing]
!> 3. [Step 3: Post-processing and cleanup]
!>
!> @see [Related routines, modules, or external references]
!> @author GSI Development Team
!> @date 2025-01-XX
!> @version [Version information if applicable]""",
            
            function_template="""!> @brief [Brief description of function purpose and return value]
!> @details [Detailed description of computation and mathematical approach]
!> [Include algorithm complexity and numerical considerations]
!>
!> @param[in] [param_name] [Description with units, constraints, and valid ranges]
!> @return [Detailed description of return value, type, units, and meaning]
!>
!> @usage
!> result = function_name(param1, param2)
!> ! [Example usage with typical parameter values and expected results]
!>
!> @algorithm
!> [Mathematical formula or computational steps]
!> [Complexity: O(n), O(n²), etc.]
!>
!> @see [Related functions, algorithms, or mathematical references]
!> @author GSI Development Team  
!> @date 2025-01-XX
!> @note [Any important notes about numerical stability, performance, etc.]""",
            
            module_template="""!> @brief [Brief description of module purpose and scope]
!> @details [Comprehensive description of module functionality and architecture]
!> [Explain the module's role in the larger GSI/EnKF system]
!>
!> This module contains the following key components:
!> - [Component 1]: [Description and purpose]
!> - [Component 2]: [Description and purpose]
!> - [Component 3]: [Description and purpose]
!>
!> Key dependencies: [List important module dependencies]
!> Performance considerations: [Memory usage, computational complexity]
!>
!> @author GSI Development Team
!> @date 2025-01-XX
!> @version [Module version information]""",
            
            parameter_tags=["@param[in]", "@param[out]", "@param[inout]"],
            return_value_tags=["@return", "@returns"],
            example_tags=["@usage", "@example"],
            math_notation="LaTeX-style notation for equations"
        )
    
    def _load_classification_templates(self) -> Dict[str, Dict[str, str]]:
        """Load classification-specific prompt templates"""
        templates = {
            "background_grid": {
                "focus_prompt": """Focus on grid operations and spatial transformations:
- Document grid types, coordinate systems, and domain decomposition
- Explain interpolation methods and parallel communication patterns
- Describe spectral transforms and boundary condition handling
- Include details about memory layout and computational efficiency""",
                "example_context": "background error covariance, grid interpolation, spectral transforms"
            },
            
            "core_analysis": {
                "focus_prompt": """Focus on data assimilation algorithms and optimization:
- Document variational methods, cost functions, and minimization algorithms
- Explain control variable transformations and preconditioning
- Describe convergence criteria and numerical stability considerations
- Include mathematical formulation of analysis equations""",
                "example_context": "cost function minimization, analysis increments, ensemble methods"
            },
            
            "io_interface": {
                "focus_prompt": """Focus on file I/O operations and data management:
- Document file format handling and data serialization
- Explain model coupling interfaces and diagnostic file generation
- Describe parallel I/O operations and error handling procedures
- Include details about metadata management and data validation""",
                "example_context": "NetCDF files, GRIB encoding, diagnostic outputs"
            },
            
            "observation_processing": {
                "focus_prompt": """Focus on observation operators and quality control:
- Document forward models and observation space transformations
- Explain quality control algorithms and bias correction methods
- Describe data thinning strategies and innovation statistics
- Include details about instrument-specific processing""",
                "example_context": "observation operators, quality control, bias correction"
            },
            
            "utilities": {
                "focus_prompt": """Focus on utility functions and support operations:
- Document mathematical operations and algorithm implementations
- Explain memory management and performance optimization
- Describe error handling and diagnostic capabilities
- Include details about numerical precision and stability""",
                "example_context": "mathematical utilities, memory management, error handling"
            }
        }
        return templates
    
    def _load_custom_prompts(self) -> Dict[str, str]:
        """Load custom prompts from external files"""
        custom_prompts = {}
        
        # Check for classification-specific prompt files
        for classification in ["background_grid", "core_analysis", "io_interface", 
                              "observation_processing", "utilities"]:
            prompt_file = self.config_dir / f"prompts_{classification}.md"
            if prompt_file.exists():
                try:
                    with open(prompt_file, 'r', encoding='utf-8') as f:
                        custom_prompts[classification] = f.read()
                    logger.info(f"Loaded custom prompts for {classification}")
                except Exception as e:
                    logger.warning(f"Failed to load custom prompts for {classification}: {e}")
        
        return custom_prompts
    
    def load_prompts(self, classification: str) -> PromptConfig:
        """Enhanced prompt loading with robust classification mapping and validation"""
        try:
            # Validate classification
            if not self._is_valid_classification(classification):
                logger.warning(f"Invalid classification '{classification}', using default")
                classification = "utilities"
            
            # Load base prompts
            base_prompt = self._load_base_prompt()
            background_context = self._load_background_context()
            
            # Use custom prompts if available, otherwise use built-in templates
            if classification in self.custom_prompts:
                classification_specific = self.custom_prompts[classification]
                logger.debug(f"Using custom prompts for {classification}")
            else:
                classification_specific = self._get_classification_prompt(classification)
            
            # Enhanced focus areas and context with validation
            focus_areas = self._get_focus_areas(classification)
            template_context = self.classification_templates.get(classification, {})
            
            # Create enhanced prompt configuration
            prompt_config = PromptConfig(
                base_prompt=base_prompt,
                classification_specific=classification_specific,
                background_context=background_context,
                focus_areas=focus_areas,
                comment_standard=self.comment_standard
            )
            
            # Validate prompt configuration completeness
            self._validate_prompt_config(prompt_config, classification)
            
            return prompt_config
            
        except Exception as e:
            logger.error(f"Error loading prompts for {classification}: {e}")
            return self._get_fallback_prompt_config(classification)
    
    def _is_valid_classification(self, classification: str) -> bool:
        """Validate if classification is supported"""
        valid_classifications = [
            "background_grid", "core_analysis", "io_interface",
            "observation_processing", "utilities"
        ]
        return classification in valid_classifications
    
    def _validate_prompt_config(self, config: PromptConfig, classification: str) -> None:
        """Validate prompt configuration completeness"""
        if not config.base_prompt.strip():
            logger.warning(f"Empty base prompt for classification: {classification}")
        
        if not config.classification_specific.strip():
            logger.warning(f"Empty classification-specific prompt for: {classification}")
            
        if not config.focus_areas:
            logger.warning(f"No focus areas defined for classification: {classification}")
        
        if not config.comment_standard:
            logger.error(f"Missing comment standard for classification: {classification}")
    
    def _get_fallback_prompt_config(self, classification: str) -> PromptConfig:
        """Get fallback prompt configuration for error cases"""
        return PromptConfig(
            base_prompt="Add comprehensive comments to this Fortran code.",
            classification_specific=f"Focus on {classification.replace('_', ' ')} functionality.",
            background_context="This is a GSI/EnKF Fortran code file.",
            focus_areas=["documentation", "clarity", "maintainability"],
            comment_standard=self.comment_standard or CommentStandard()
        )
    
    def _load_base_prompt(self) -> str:
        """Load base prompt from prompts.md file"""
        prompts_file = self.config_dir / "prompts.md"
        if prompts_file.exists():
            try:
                with open(prompts_file, 'r', encoding='utf-8') as f:
                    return f.read()
            except Exception as e:
                logger.error(f"Error loading base prompt: {e}")
        return ""
    
    def _load_background_context(self) -> str:
        """Load background context from background.md file"""
        background_file = self.config_dir / "background.md"
        if background_file.exists():
            try:
                with open(background_file, 'r', encoding='utf-8') as f:
                    return f.read()
            except Exception as e:
                logger.error(f"Error loading background context: {e}")
        return ""
    
    def create_custom_prompt_template(self, classification: str, template_content: str) -> bool:
        """Create a custom prompt template file for a classification"""
        try:
            prompt_file = self.config_dir / f"prompts_{classification}.md"
            with open(prompt_file, 'w', encoding='utf-8') as f:
                f.write(template_content)
            
            # Reload custom prompts to include the new template
            self.custom_prompts = self._load_custom_prompts()
            logger.info(f"Created custom prompt template for {classification}")
            return True
            
        except Exception as e:
            logger.error(f"Error creating custom prompt template for {classification}: {e}")
            return False
    
    def validate_prompt_configuration(self) -> Dict[str, Any]:
        """Validate prompt configuration and return status"""
        validation_results = {
            "base_prompt_exists": False,
            "background_context_exists": False,
            "custom_prompts_count": 0,
            "template_completeness": {},
            "errors": []
        }
        
        try:
            # Check base files
            validation_results["base_prompt_exists"] = (self.config_dir / "prompts.md").exists()
            validation_results["background_context_exists"] = (self.config_dir / "background.md").exists()
            validation_results["custom_prompts_count"] = len(self.custom_prompts)
            
            # Check template completeness for each classification
            for classification in ["background_grid", "core_analysis", "io_interface", 
                                  "observation_processing", "utilities"]:
                config = self.load_prompts(classification)
                validation_results["template_completeness"][classification] = {
                    "has_base_prompt": bool(config.base_prompt.strip()),
                    "has_classification_specific": bool(config.classification_specific.strip()),
                    "has_background_context": bool(config.background_context.strip()),
                    "has_comment_standard": config.comment_standard is not None,
                    "focus_areas_count": len(config.focus_areas)
                }
                
        except Exception as e:
            validation_results["errors"].append(f"Validation error: {e}")
            
        return validation_results
    
    def _get_classification_prompt(self, classification: str) -> str:
        """Get classification-specific prompt"""
        prompts = {
            "background_grid": "Focus on grid operations, domain decomposition, coordinate transformations, and interpolation routines.",
            "core_analysis": "Focus on analysis algorithms, cost function computation, minimization routines, and control vector operations.", 
            "io_interface": "Focus on file I/O operations, data formats, NetCDF/GRIB interfaces, and diagnostic file generation.",
            "observation_processing": "Focus on observation operators, quality control, data thinning, and measurement processing.",
            "utilities": "Focus on utility functions, mathematical operations, helper routines, and support functions."
        }
        return prompts.get(classification, "Focus on general Fortran code functionality and structure.")
    
    def _get_focus_areas(self, classification: str) -> List[str]:
        """Get focus areas for classification"""
        focus_map = {
            "background_grid": ["grid_operations", "interpolation", "coordinate_transforms"],
            "core_analysis": ["cost_functions", "minimization", "control_vectors"],
            "io_interface": ["file_io", "data_formats", "diagnostics"],
            "observation_processing": ["obs_operators", "quality_control", "data_processing"],
            "utilities": ["mathematical_ops", "helper_functions", "support_routines"]
        }
        return focus_map.get(classification, ["general_functionality"])